import pandas as pd
from gensim.models.doc2vec import Doc2Vec, TaggedDocument

def getFeat():
    train = pd.read_csv("Train.csv", encoding = "ISO-8859-1")
    train_words = [word.split() for word in train['text']]
    test = pd.read_csv("Test.csv", encoding = "ISO-8859-1")
    test_words = [word.split() for word in test['text']]
    words = train_words+test_words

    documents = [TaggedDocument(doc, [i]) for i, doc in enumerate(words)]
    dv_model = Doc2Vec(documents, vector_size=10, window=2, min_count=1, workers=4)

    matrix_dv = []
    matrix_test = []
    i = 0
    for sentence, _ in documents:
        if(i<len(train_words)):
            matrix_dv.append(dv_model.infer_vector(sentence))
        else:
            matrix_test.append(dv_model.infer_vector(sentence))
    output = []
    if 'labels' in train.columns:
        for label in train['labels']:
            if(label == 'positive'):
                output.append(0)
            elif(label == 'neutral'):
                output.append(1)
            else:
                output.append(2)
    return matrix_dv,output,matrix_test

